import torch
import torchtext; 
torchtext.disable_torchtext_deprecation_warning()
import pickle
import torch.nn as nn
from utils import list2sentence, tensor2sentence, translate_sentence, save_checkpoint, load_checkpoint
from torch.utils.data import DataLoader
from dataset import NumberDataset
from model import Transformer
from sklearn.model_selection import train_test_split


# hyperparameters 1

BATCH_SIZE = 32
EPOCHS = 10000
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
learning_rate = 3e-4
max_len = 120


# get vocabulary

zh_vocab_file_path = "zh_vocab.pkl"
en_vocab_file_path = "en_vocab.pkl"

with open(zh_vocab_file_path, "rb") as zh_vocab_file:
    zh_vocab = pickle.load(zh_vocab_file)

with open(en_vocab_file_path, "rb") as en_vocab_file:
    en_vocab = pickle.load(en_vocab_file)

zh_ivocab = {index: token for token, index in zh_vocab.items()}
en_ivocab = {index: token for token, index in en_vocab.items()}

# hyperparameters 2(for model)

src_vocab_size = len(zh_vocab)
trg_vocab_size = len(en_vocab)
embedding_size = 512
num_heads = 8
num_layers = 3
num_encoder_layers = 3
num_decoder_layers = 3
dropout = 0.10
forward_expansion = 4
src_pad_idx = 0
trg_pad_idx = 0


# load model

model = Transformer(
    src_vocab_size,
    trg_vocab_size,
    src_pad_idx,
    trg_pad_idx,
    embedding_size,
    num_layers,
    forward_expansion,
    num_heads,
    dropout,
    DEVICE,
    max_len
).to(DEVICE)

'''
model = Transformer(
    embedding_size,
    src_vocab_size,
    trg_vocab_size,
    src_pad_idx,
    num_heads,
    num_encoder_layers,
    num_decoder_layers,
    forward_expansion,
    dropout,
    max_len,
    DEVICE,
).to(DEVICE)
'''


optimizer = torch.optim.Adam(
    model.parameters(), lr=learning_rate
)


load_checkpoint(torch.load("my_checkpoint.pth.tar"), model, optimizer)   

#sentence = "经济危机的不断加深使我们看到了危机过后的世界是什么样子的。"

while True:

    sentence = input("")
    print("Q:",sentence)
    translated_sentence = translate_sentence(model, sentence,  zh_vocab, en_ivocab, DEVICE, max_len=50)
    print(translated_sentence)
